{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "473b0cb3",
   "metadata": {},
   "source": [
    "## The Monty Hall Problem\n",
    "\n",
    "The Monty Hall problem arose from the gameshow <i>Let's Make a Deal</i>, where a guest had to choose which one of three doors had a prize behind it. The twist was that after the guest chose, the host, originally Monty Hall, would then open one of the doors the guest did not pick that also did not have the prize behind it. Afterwards, Monty would ask if the guest wanted to switch which door they had picked. Initial inspection may lead you to believe that if there are only two doors left there is a 50-50 chance of you picking the right one, and so there is no advantage one way or the other. However, it has been proven both through simulations and analytically that there is in fact a 66% chance of getting the prize if the guest switches their door after Monty opens one, regardless of the door they initially went with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "76bedfe3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Populating the interactive namespace from numpy and matplotlib\n",
      "torch        : 1.13.0\n",
      "torchegranate: 0.4.0\n",
      "\n",
      "Compiler    : GCC 11.2.0\n",
      "OS          : Linux\n",
      "Release     : 4.15.0-206-generic\n",
      "Machine     : x86_64\n",
      "Processor   : x86_64\n",
      "CPU cores   : 8\n",
      "Architecture: 64bit\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%pylab inline\n",
    "import seaborn; seaborn.set_style('whitegrid')\n",
    "import torch\n",
    "\n",
    "%load_ext watermark\n",
    "%watermark -m -n -p torch,torchegranate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dccea1d",
   "metadata": {},
   "source": [
    "We can reproduce this result in pomegranate using Bayesian networks with three nodes, one for the guest, one for the prize, and one for the door Monty chooses to open. The door the guest initially chooses and the door the prize is behind are completely random processes across the three doors, but the door which Monty opens is dependent on both the door the guest chooses (it cannot be the door the guest chooses), and the door the prize is behind (it cannot be the door with the prize behind it).\n",
    "\n",
    "To create the Bayesian network in pomegranate, we first create the distributions which live in each node in the graph. For a categorical bayesian network we use Categorical distributions for the root nodes and ConditionalCategorical distributions for the inner and leaf nodes. \n",
    "\n",
    "First, we can create our \"prize\" and \"guest\" distributions. These are each Categorical distributions because they do not depend on anything, and they are uniform distributions because they are equally likely to be any of the three doors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ce8a68c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchegranate.distributions import Categorical\n",
    "\n",
    "guest = Categorical([[1./3, 1./3, 1./3]])\n",
    "prize = Categorical([[1./3, 1./3, 1./3]])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "369339f5",
   "metadata": {},
   "source": [
    "You may notice that there is an additional dimension added to the probabilities. This is because all distributions in pomegranate have the potential to be multivariate even when being applied to univariate problems.\n",
    "\n",
    "Next, we need to create the conditional distribution describing the door that Monty will open. Because Monty can only open a door that is not selected by the contestant and also does not have the prize, sometimes this leaves Monty with only one door that can be opened. Overall, the distribution is a 3x3x3 tensor, with three possibilities from the guest, three independent possibilities from the prize, and three possible doors to open."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2b303ce6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchegranate.distributions import ConditionalCategorical\n",
    "\n",
    "probs = numpy.array([[\n",
    "     [[0.0, 0.5, 0.5], [0.0, 0.0, 1.0], [0.0, 1.0, 0.0]], \n",
    "     [[0.0, 0.0, 1.0], [0.5, 0.0, 0.5], [1.0, 0.0, 0.0]],\n",
    "     [[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.5, 0.5, 0.0]]\n",
    "]])\n",
    "\n",
    "monty = ConditionalCategorical(probs) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f693f22a",
   "metadata": {},
   "source": [
    "Next, we can create the Bayesian network object in just one line by passing in the distribution objects and edges in the form of (parent, child) tuples. Previous versions of pomegranate required that you create State or Node objects and add them in using `add_edge` and `add_node` methods. State and Node objects no longer exist, and while those methods still exist if you would prefer to use them you no longer need to. The `bake` method has also been removed and is no longer required."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cdd24ba0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchegranate.bayesian_network import BayesianNetwork\n",
    "\n",
    "model = BayesianNetwork([guest, prize, monty], [(guest, monty), (prize, monty)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ec59baca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0, 1, 2],\n",
       "        [0, 2, 1],\n",
       "        [2, 1, 0]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = torch.tensor([[0, 1, -1],\n",
    "                  [0, 2, -1],\n",
    "                  [2, 1, -1]])\n",
    "\n",
    "X_masked = torch.masked.MaskedTensor(X, mask=X >= 0)\n",
    "\n",
    "\n",
    "model.predict(X_masked)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "dac5d471",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([1.6111])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torchegranate.distributions import Exponential\n",
    "\n",
    "X = torch.exp(torch.randn(100, 1))\n",
    "mask = torch.ones(100, 1, dtype=bool)\n",
    "mask[75:] = False\n",
    "X_masked = torch.masked.MaskedTensor(X, mask=mask)\n",
    "\n",
    "Exponential().fit(X[:75]).scales"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "6169a52f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([1.6111])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Exponential().fit(X_masked).scales"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ea963c8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
